最低限の SageMaker HyperPod クラスターを HashiCorp Terraform で作ってみた

#SageMaker HyperPod
#AWS
たかくに
2024.12.13
こんにちは！AWS 事業本部コンサルティング部のたかくに（@takakuni_）です。
Amazon SageMaker HyperPod 使っていますか？
今回は普段の検証用に作ったり削除できたりできるよう、 HashiCorp Terraform で SageMaker HyperPod を作れるようにしておきます。
コードは以下に格納されていますので、全体像をご覧になりたい場合は以下をご覧ください。
https://github.com/takakuni-classmethod/genai-blog/tree/main/sagemaker_hyperpod_101
 構成今回は非常にシンプルな最低限の SageMaker HyperPod クラスターを作成します。
オーケストレータは Slurm で、 FSx for Lustre 等のファイルシステムは作成しません。
また、SageMaker 側で提供される VPC 内に SageMaker HyperPod クラスターを作成します。非常にシンプルですね。
 やってみるそれでは早速作成してきましょう。
執筆時点で最新の AWS Provider v5.81.0 は SageMaker HyperPod クラスターをサポートしていません。
そのため、現時点では AWS Cloud Control Provider v1.23.0 を併用して作成します。
providers.tf
terraform {
  required_providers {
    aws = {
      source  = "hashicorp/aws"
      version = "5.81.0"
    }
    awscc = {
      source  = "hashicorp/awscc"
      version = "1.23.0"
    }
  }
}

provider "aws" {
  # Configuration options
  region = "ap-northeast-1"
}

provider "awscc" {
  # Configuration options
  region = "ap-northeast-1"
}

data "aws_caller_identity" "self" {}
data "aws_region" "current" {}

locals {
  account_id = data.aws_caller_identity.self.account_id
  region     = data.aws_region.current.name
}
 IAM RoleSageMaker HyperPod のインスタンスグループに紐づける実行ロールを作成します。
ユーザーの VPC に所属しない場合は、管理ポリシーの AmazonSageMakerClusterInstanceRolePolicy で事足ります。非常に簡単ですね。
https://docs.aws.amazon.com/aws-managed-policy/latest/reference/AmazonSageMakerClusterInstanceRolePolicy.html
resource "aws_iam_role" "hyperpod" {
  name = "sagemaker-hyperpod-role"
  assume_role_policy = jsonencode({
    Version = "2012-10-17",
    Statement = [
      {
        Effect = "Allow",
        Principal = {
          Service = "sagemaker.amazonaws.com"
        },
        Action = "sts:AssumeRole"
      }
    ]
  })
}

resource "aws_iam_role_policy_attachment" "hyperpod_managed" {
  role       = aws_iam_role.hyperpod.name
  policy_arn = "arn:aws:iam::aws:policy/AmazonSageMakerClusterInstanceRolePolicy"
}
ユーザー VPC 内に所属させる場合は以下の権限が追加で必要になります。
https://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/sagemaker-hyperpod-prerequisites-iam.html#sagemaker-hyperpod-prerequisites-iam-role-for-hyperpod
 S3 の作成ノード起動時にセットアップするためのスクリプトを配置する S3 バケットを作成します。
ワイルドカードと for_each で、フォルダ配下のスクリプトを一括アップロードできるのは Terraform の素晴らしいポイントですね。
s3.tf
resource "aws_s3_bucket" "life_cycle_scripts" {
  bucket = "sagemaker-hyperpod-lifecycle-${local.account_id}"
}

resource "aws_s3_bucket_public_access_block" "life_cycle_scripts" {
  bucket                  = aws_s3_bucket.life_cycle_scripts.bucket
  block_public_acls       = true
  block_public_policy     = true
  ignore_public_acls      = true
  restrict_public_buckets = true
}

resource "aws_s3_bucket_ownership_controls" "life_cycle_scripts" {
  bucket = aws_s3_bucket.life_cycle_scripts.bucket
  rule {
    object_ownership = "BucketOwnerPreferred"
  }
}

resource "aws_s3_object" "life_cycle_scripts" {
  for_each = fileset("./config/", "**")
  bucket   = aws_s3_bucket.life_cycle_scripts.bucket
  key      = "config/${each.value}"
  source   = "./config/${each.value}"
}

resource "aws_s3_object" "life_cycle_scripts_provisioning_parameters" {
  bucket = aws_s3_bucket.life_cycle_scripts.bucket
  key    = "config/provisioning_parameters.json"

  content = jsonencode(
    {
      "version"          = "1.0.0"
      "workload_manager" = "slurm"
      "controller_group" = "controller-machine"
      "worker_groups" = [
        {
          "instance_group_name" = "worker-group-1"
          "instance_type"       = "ml.t3.medium"
        }
      ]
    }
  )
}
config ファイルは以下から拝借しました。
https://github.com/aws-samples/awsome-distributed-training/tree/main/1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config
 SageMaker HyperPod クラスター最後に SageMaker HyperPod クラスターを作成します。
ml.t3.medium から動くので料金も抑えめで動かします。
また、今回はコントローラー側のみ起動します。
hyperpod.tf
resource "awscc_sagemaker_cluster" "this" {
  cluster_name = "sagemaker-hyperpod"

  instance_groups = [{
    execution_role      = aws_iam_role.hyperpod.arn
    instance_count      = 1
    instance_group_name = "controller-machine"
    instance_type       = "ml.t3.medium"
    life_cycle_config = {
      source_s3_uri = "s3://${aws_s3_bucket.life_cycle_scripts.id}/config/"
      on_create     = "on_create.sh"
    }
  }]
}
terraform apply　を行い、10 分ほど経過しました。コンソールを確認するとうまく動いてそうです。
 ノードに接続作成したコントローラーノードにログインします。
CloudShell から Systems Manager でログインしてみます。
ターゲットは sagemaker-cluster:<cluster_id>_<node-group>-<instance_id> の命名です。
aws ssm start-session \
    --target sagemaker-cluster:xfcguk1uawxc_controller-machine-i-08f42f748b8ec24d1 \
    --region ap-northeast-1
ログインできていますね。
[cloudshell-user@ip-10-132-64-9 ~]$ aws ssm start-session \
>     --target sagemaker-cluster:xfcguk1uawxc_controller-machine-i-08f42f748b8ec24d1 \
>     --region ap-northeast-1

Starting session with SessionId: cm-takakuni-vy6bvhzcn75yggsf2ten2xebt8
# whoami
root
# sinfo
PARTITION AVAIL  TIMELIMIT  NODES  STATE NODELIST
dev*         up   infinite      0    n/a
#
 まとめ以上、「最低限の SageMaker HyperPod クラスターを HashiCorp Terraform で作ってみた」でした。
大規模言語モデルを作成するために、インスタンスモリモリなケースが多いですが、使用感を確かめるためにも ml.t3.medium など小さいサイズのインスタンスタイプも用意されています。
このブログがどなたかの参考になれば幸いです。
AWS 事業本部コンサルティング部のたかくに（@takakuni_）でした！